import cv2
import numpy as np
 
def drawMatchesKnn_cv2(img1, kp1, img2, kp2, goodMatch):
    h1, w1 = img1.shape[:2]
    h2, w2 = img2.shape[:2]
 
    vis = np.zeros((max(h1, h2), w1 + w2, 3), np.uint8) #TODO
    # vis = np.zeros((max(h1, h2), w1 + w2), np.uint8)
    vis[:h1, :w1] = img1
    vis[:h2, w1:w1 + w2] = img2
 
    p1 = [kpp.queryIdx for kpp in goodMatch]
    p2 = [kpp.trainIdx for kpp in goodMatch]
 
    post1 = np.int32([kp1[pp].pt for pp in p1])
    post2 = np.int32([kp2[pp].pt for pp in p2]) + (w1, 0)
 
    for (x1, y1), (x2, y2) in zip(post1, post2):
        cv2.line(vis, (x1, y1), (x2, y2), (0,0,255), thickness=5)
 
    cv2.namedWindow("match", cv2.WINDOW_NORMAL)
    cv2.imshow("match", vis)
    cv2.imwrite("data_out/out.jpg", vis)

img1 = cv2.imread("data_in/axle-3.jpg", cv2.IMREAD_COLOR)
img2 = cv2.imread("data_in/3.bmp", cv2.IMREAD_COLOR)

img1 = cv2.resize(img1, (1280, 4210))   # w, h
img2 = cv2.resize(img2, (1280, 4210))

img1_1 = img1[(851 - 640):(851 + 640), :] # h, w
img1_2 = img1[(3351 - 640):(3351 + 640), :]
img2_1 = img2[(851 - 640):(851 + 640), :]
img2_2 = img2[(3351 - 640):(3351 + 640), :]

cv2.imwrite("data_out/img1_1.jpg", img1_1)
cv2.imwrite("data_out/img1_2.jpg", img1_2)
cv2.imwrite("data_out/img2_1.jpg", img2_1)
cv2.imwrite("data_out/img2_2.jpg", img2_2)

img1_1 = cv2.imread("data_in/img2_1_normal.jpg", cv2.IMREAD_COLOR)

# img1 = cv2.cvtColor(img1, cv2.COLOR_BGR2GRAY)
# img2 = cv2.cvtColor(img2, cv2.COLOR_BGR2GRAY)

feature = cv2.xfeatures2d.SIFT_create()
# feature = cv2.xfeatures2d.SURF_create()
# feature = cv2.ORB_create()

kp1, des1 = feature.detectAndCompute(img1_1, None)
kp2, des2 = feature.detectAndCompute(img2_1, None)

# BFmatcher with default parms
bf = cv2.BFMatcher(cv2.NORM_L2)
matches = bf.knnMatch(des1, des2, k = 2)
 
goodMatch = []
for m,n in matches:
    if m.distance < 0.60 * n.distance:
        goodMatch.append(m)
print('length of good match: {}'.format(len(goodMatch)))

if len(goodMatch) > 0:
    drawMatchesKnn_cv2(img1_1, kp1, img2_1, kp2, goodMatch[:20])

    cv2.waitKey(0)
    cv2.destroyAllWindows()